package gpucore

import (
	"fmt"
	"log"
	"runtime"

	"github.com/cogentcore/webgpu/wgpu"
)

// GPUContext 是 WebGPU 的高层抽象
type GPUContext interface {
	CreateBuffer(data interface{}, usage BufferUsage) (GPUBuffer, error)
	RegisterKernel(name string, wgsl string, layout *wgpu.BindGroupLayout, workgroupSize [3]uint32) error
	GetKernel(name string) (ComputeKernel, bool)
	Wait() error
	Close()
	Device() *wgpu.Device
}

type gpuContext struct {
	instance *wgpu.Instance
	adapter  *wgpu.Adapter
	device   *wgpu.Device
	queue    *wgpu.Queue
	kernels  map[string]*compiledKernel
}

// NewGPUContext 初始化 WebGPU 上下文
func NewGPUContext() (GPUContext, error) {
	wgpu.SetLogLevel(wgpu.LogLevelInfo)
	log.Println("Initializing WebGPU context...")

	instance := wgpu.CreateInstance(nil)
	if instance == nil {
		return nil, fmt.Errorf("failed to create WebGPU instance")
	}

	adapter, err := instance.RequestAdapter(&wgpu.RequestAdapterOptions{
		PowerPreference: wgpu.PowerPreferenceHighPerformance,
	})
	if err != nil {
		instance.Release()
		return nil, fmt.Errorf("request adapter: %w", err)
	}

	device, err := adapter.RequestDevice(&wgpu.DeviceDescriptor{})
	if err != nil {
		adapter.Release()
		instance.Release()
		return nil, fmt.Errorf("request device: %w", err)
	}

	ctx := &gpuContext{
		instance: instance,
		adapter:  adapter,
		device:   device,
		queue:    device.GetQueue(),
		kernels:  make(map[string]*compiledKernel),
	}

	runtime.SetFinalizer(ctx, func(c *gpuContext) { c.Close() })
	return ctx, nil
}

func (c *gpuContext) RegisterKernel(name string, wgsl string, layout *wgpu.BindGroupLayout, wgs [3]uint32) error {
	if layout == nil {
		return fmt.Errorf("bind group layout for kernel '%s' is nil", name)
	}

	module, err := c.device.CreateShaderModule(&wgpu.ShaderModuleDescriptor{
		WGSLDescriptor: &wgpu.ShaderModuleWGSLDescriptor{Code: wgsl},
	})
	if err != nil {
		return fmt.Errorf("compile shader '%s': %w", name, err)
	}

	pipelineLayout, err := c.device.CreatePipelineLayout(&wgpu.PipelineLayoutDescriptor{
		BindGroupLayouts: []*wgpu.BindGroupLayout{layout},
	})
	if err != nil {
		module.Release()
		return fmt.Errorf("create pipeline layout for '%s': %w", name, err)
	}

	pipeline, err := c.device.CreateComputePipeline(&wgpu.ComputePipelineDescriptor{
		Layout: pipelineLayout,
		Compute: wgpu.ProgrammableStageDescriptor{
			Module:     module,
			EntryPoint: "main",
		},
	})
	if err != nil {
		pipelineLayout.Release()
		module.Release()
		return fmt.Errorf("create compute pipeline for '%s': %w", name, err)
	}

	pipelineLayout.Release()
	module.Release()

	c.kernels[name] = &compiledKernel{
		device:        c.device,
		pipeline:      pipeline,
		bindLayout:    layout,
		workgroupSize: wgs,
	}
	return nil
}

func (c *gpuContext) GetKernel(name string) (ComputeKernel, bool) {
	k, ok := c.kernels[name]
	return k, ok
}

func (c *gpuContext) CreateBuffer(data interface{}, usage BufferUsage) (GPUBuffer, error) {
	wgpuUsage := usage.toWGPU()
	var buf *wgpu.Buffer
	var size uint64

	switch v := data.(type) {
	case []float32:
		bytes := wgpu.ToBytes(v)
		size = uint64(len(bytes))
		var err error
		buf, err = c.device.CreateBufferInit(&wgpu.BufferInitDescriptor{
			Contents: bytes,
			Usage:    wgpuUsage,
		})
		if err != nil {
			return nil, fmt.Errorf("create initialized buffer: %w", err)
		}
	case []uint32:
		bytes := wgpu.ToBytes(v)
		size = uint64(len(bytes))
		var err error
		buf, err = c.device.CreateBufferInit(&wgpu.BufferInitDescriptor{
			Contents: bytes,
			Usage:    wgpuUsage,
		})
		if err != nil {
			return nil, fmt.Errorf("create initialized uint32 buffer: %w", err)
		}
	default:
		return nil, fmt.Errorf("unsupported data type: %T (only []float32, []uint32 supported)", data)
	}

	return newGPUBuffer(c.device, buf, size), nil
}

func (c *gpuContext) Wait() error {
	c.device.Poll(true, nil)
	return nil
}

func (c *gpuContext) Close() {
	for _, k := range c.kernels {
		if k.pipeline != nil {
			k.pipeline.Release()
		}
		if k.bindLayout != nil {
			k.bindLayout.Release()
		}
	}
	c.kernels = nil

	if c.device != nil {
		c.device.Release()
		c.device = nil
	}
	if c.adapter != nil {
		c.adapter.Release()
		c.adapter = nil
	}
	if c.instance != nil {
		c.instance.Release()
		c.instance = nil
	}
}
func (c *gpuContext) Device() *wgpu.Device {
	return c.device
}
